/**
 * @file bufferPool.hpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-09
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */

#pragma once

#include <malloc.h>
#include <cassert>
#include <condition_variable>
#include <iostream>
#include <mutex>
#include <queue>
#include "logging.h"
#include "lyn_memory.h"
#include "util.hpp"

/********************************************************************************
类名 : BufferPool
Description: 内存池模板类
Example: BufferPool oPool(100, 10, DEVICE)
*******************************************************************************/

using namespace std::chrono;

/**
 * @brief 内存所属
 */
typedef enum {
  DEVICE = 0,  // 所属设备侧
  HOST         // 所属系统侧
} MEM_SIDE_E;

class BufferPool {
 public:
  /**
   * @brief 构造函数，根据类型来决定申请主机侧还是设备侧内存
   *
   * @param[in] elementSize 表示单个内存元素大小
   * @param[in] elementCount 表示内存池容量
   * @param[in] side 表示内存所属
   */
  BufferPool(uint32_t elementSize, uint32_t elementCount,
             MEM_SIDE_E side = DEVICE) {
    m_elementSize = elementSize;
    m_elementCount = elementCount;
    m_memSide = side;
    for (uint32_t i = 0; i < m_elementCount; i++) {
      void *buffer = nullptr;
      if (m_memSide == DEVICE) {
        CHECK_ERR(lynMalloc(&buffer, m_elementSize));
      } else {
        buffer = malloc(m_elementSize);
      }
      if (buffer == nullptr) {
        loge("buffer is nullptr!");
        quick_exit(-1);
      }
      m_queue.push(buffer);
    }
  }

  ~BufferPool() {
    {
      std::unique_lock<std::mutex> lock(m_mutex);
      if (!m_condVar.wait_for(lock, std::chrono::seconds(60), [this] {
            return m_queue.size() == m_elementCount;
          })) {
        std::cout << "[Warning]: m_queue is not full, some element may lose "
                     "somewhere outside."
                  << std::endl;
      }
    }

    void *buffer;
    uint32_t size = m_queue.size();
    for (uint32_t i = 0; i < size; i++) {
      buffer = m_queue.front();
      if (m_memSide == DEVICE) {
        lynFree(buffer);
      } else {
        free(buffer);
      }
      m_queue.pop();
    }
  }

  /**
   * @brief 将内存归还内存池
   *
   * @param[in] buffer 归还的内存地址
   * @return 无
   */
  void Push(void *buffer) {
    if (Full()) {
      loge("Pool is full, but buffer still push.");
      quick_exit(-1);
    }

    std::lock_guard<std::mutex> lock(m_mutex);
    m_queue.push(buffer);
    m_condVar.notify_one();
  }

  /**
   * @brief 从内存池取出内存，1分钟超时时间内没有可用内存时，返回nullptr
   *
   * @return 内存地址
   */
  void *Pop() {
    std::unique_lock<std::mutex> lock(m_mutex);
    if (!m_condVar.wait_for(lock, std::chrono::seconds(60),
                            [this] { return !m_queue.empty(); })) {
      std::cout << "[Error]: m_queue is empty, pop null element." << std::endl;
      return nullptr;
    }

    void *buffer = m_queue.front();
    m_queue.pop();
    return buffer;
  }

  /**
   * @brief 检查内存池可用数量是否到达容量
   *
   * @return true表示已达到，false表示未达到
   */
  bool Full() {
    std::lock_guard<std::mutex> lock(m_mutex);
    return (m_queue.size() == m_elementCount);
  }

  /**
   * @brief 获取内存池可用数量
   *
   * @return 返回内存池可用数量
   */
  uint32_t Size() {
    std::lock_guard<std::mutex> lock(m_mutex);
    return m_queue.size();
  }

  /**
   * @brief 获取单个内存元素大小
   *
   * @return 返回单个内存元素大小
   */
  uint32_t ElementSize() { return m_elementSize; }

 protected:
  std::mutex m_mutex;
  std::condition_variable m_condVar;
  std::queue<void *> m_queue;

  uint32_t m_elementSize;   // 单个内存元素大小
  uint32_t m_elementCount;  // 内存池容量
  MEM_SIDE_E m_memSide;     // 内存所属
};                          // BufferPool
